import cv2
import numpy as np
import math
import random
import json
import os

def load_label(label_path):
    labels = []
    with open(label_path, 'r') as fin:
        label_obj = json.load(fin)
        img_h = label_obj['imageHeight']
        img_w = label_obj['imageWidth']

        #assert len(label_obj['shapes'] == 1)
        for obj in label_obj['shapes']:
            label = obj['label']
            point = obj['points']
            left_top = point[0]
            right_bottom = point[1]
            left, top = map(float,left_top)
            right, bottom = map(float, right_bottom)
            w = right - left
            h = bottom - top
            x_center = left + w/2
            y_center = top + h/2
            x_c = x_center/img_w
            y_c = y_center/img_h
            w_ratio = w/img_w
            h_ratio = h/img_h
            labels.append((int(x_center), int(y_center), w, h, x_c, y_c, w_ratio, h_ratio))
    print(labels)
    return labels 

        

def load_background_image(background_path, target_size=(640, 480)):
    background = cv2.imread(background_path)
    if background is None:
        raise ValueError(f"无法加载背景图片: {background_path}")
    background = cv2.resize(background, target_size)
    return background

def load_basketball_texture(texture_path, size=256):
    texture = cv2.imread(texture_path, cv2.IMREAD_UNCHANGED)
    if texture is None:
        raise ValueError(f"无法加载纹理图片:{texture_path}")
    texture = cv2.resize(texture, (size, size))
    if texture.shape[2] == 3:
        alpha = np.ones((size, size, 1), dtype=np.uint8)*255
        texture = np.concatenate([texture, alpha], axis=2)
    # 创建圆形遮罩
    center = size // 2
    radius = center - 28
    mask = np.zeros((size, size), dtype=np.uint8)
    cv2.circle(mask, (center, center), radius, 255, -1)
    
    # 应用圆形遮罩到alpha通道
    texture[:, :, 3] = mask
    return texture

def create_spherical_basketball_texture(size=256):
    """创建带光照效果的篮球纹理"""
    texture = np.zeros((size, size, 4), dtype=np.uint8)  # RGBA
    
    center = size // 2
    radius = center - 2
    
    # 基础颜色 - 橙褐色
    base_color = np.array([200, 100, 50, 255])
    line_color = np.array([0, 0, 0, 255])
    
    # 预计算距离和角度 - 使用 meshgrid 而不是 ogrid
    y_coords, x_coords = np.mgrid[:size, :size]
    dx = x_coords - center
    dy = y_coords - center
    distance_sq = dx**2 + dy**2
    distance = np.sqrt(distance_sq)
    
    # 创建掩码
    mask = distance <= radius
    valid_mask = mask & (distance > 0)
    
    # 计算法线（避免数学域错误）
    nx = np.zeros_like(distance)
    ny = np.zeros_like(distance)
    nz = np.zeros_like(distance)
    
    nx[valid_mask] = dx[valid_mask] / distance[valid_mask]
    ny[valid_mask] = dy[valid_mask] / distance[valid_mask]
    
    # 安全计算nz，避免负数平方根
    nz_sq = 1 - nx**2 - ny**2
    nz = np.sqrt(np.maximum(nz_sq, 0))  # 确保非负
    
    # 中心点特殊处理
    nx[center, center] = 0
    ny[center, center] = 0
    nz[center, center] = 1
    
    # 光照计算（光源在左上角）
    light_dir = np.array([-0.7, -0.7, 0.7])
    light_dir = light_dir / np.linalg.norm(light_dir)
    
    # 计算光照强度
    intensity = nx * light_dir[0] + ny * light_dir[1] + nz * light_dir[2]
    intensity = np.maximum(intensity, 0.3)  # 确保最小亮度
    
    # 应用基础颜色和光照
    for c in range(3):
        texture[:, :, c] = (base_color[c] * intensity).astype(np.uint8)
    texture[:, :, 3] = mask.astype(np.uint8) * 255  # Alpha通道
    
    # 添加篮球线条
    angle = np.arctan2(dy, dx)
    
    # 边缘黑线
    edge_mask = (distance > radius * 0.9) & mask
    texture[edge_mask] = [*line_color[:3], 255]
    
    # 径向线条
    for i in range(12):  # 12条径向线
        line_angle = i * np.pi / 6
        line_mask = (np.abs(np.mod(angle - line_angle + np.pi, 2*np.pi) - np.pi) < 0.05) & mask
        texture[line_mask] = [*line_color[:3], 255]
    
    return texture

def apply_spherical_lighting_safe(ball_region, center, radius):
    """安全地应用球面光照效果"""
    height, width = ball_region.shape[:2]
    lit_ball = ball_region.copy()
    
    # 预计算坐标 - 使用 meshgrid
    y_coords, x_coords = np.mgrid[:height, :width]
    dx = x_coords - center[0]
    dy = y_coords - center[1]
    distance_sq = dx**2 + dy**2
    distance = np.sqrt(distance_sq)
    
    # 创建有效掩码
    valid_mask = (distance <= radius) & (distance > 0)
    center_mask = (distance == 0)
    
    # 计算法线
    nx, ny, nz = np.zeros_like(distance), np.zeros_like(distance), np.zeros_like(distance)
    
    nx[valid_mask] = dx[valid_mask] / distance[valid_mask]
    ny[valid_mask] = dy[valid_mask] / distance[valid_mask]
    
    nz_sq = 1 - nx**2 - ny**2
    nz = np.sqrt(np.maximum(nz_sq, 0))
    
    # 中心点处理
    nx[center_mask] = 0
    ny[center_mask] = 0
    nz[center_mask] = 1
    
    # 光照方向
    light_dir = np.array([-0.7, -0.7, 0.7])
    light_dir = light_dir / np.linalg.norm(light_dir)
    
    # 计算光照强度
    intensity = nx * light_dir[0] + ny * light_dir[1] + nz * light_dir[2]
    intensity = np.maximum(intensity, 0.2)
    intensity = 1.0
    
    # 应用光照
    alpha_mask = lit_ball[:, :, 3] > 0
    for c in range(3):
        lit_ball[:, :, c] = np.where(
            alpha_mask,
            np.clip(lit_ball[:, :, c] * intensity, 0, 255),
            lit_ball[:, :, c]
        )
    
    return lit_ball

def blend_basketball_naturally(background, ball_texture, position, radius):
    """自然融合篮球到背景中"""
    bg_height, bg_width = background.shape[:2]
    x, y = position
    
    # 调整篮球大小
    ball_resized = cv2.resize(ball_texture, (radius*2, radius*2))
    
    # 计算有效区域
    ball_x1 = max(0, x - radius)
    ball_x2 = min(bg_width, x + radius)
    ball_y1 = max(0, y - radius)
    ball_y2 = min(bg_height, y + radius)
    
    tex_x1 = max(0, radius - x) if x < radius else 0
    tex_x2 = min(radius*2, radius + (bg_width - x)) if x + radius > bg_width else radius*2
    tex_y1 = max(0, radius - y) if y < radius else 0
    tex_y2 = min(radius*2, radius + (bg_height - y)) if y + radius > bg_height else radius*2
    
    if tex_x1 >= tex_x2 or tex_y1 >= tex_y2:
        return background, None
    
    # 提取区域
    ball_region = ball_resized[tex_y1:tex_y2, tex_x1:tex_x2]
    bg_region = background[ball_y1:ball_y2, ball_x1:ball_x2].copy()
    
    # 应用光照
    ball_center = (ball_region.shape[1]//2, ball_region.shape[0]//2)
    ball_region_lit = apply_spherical_lighting_safe(ball_region, ball_center, min(ball_center[0], ball_center[1]))
    
    # Alpha混合
    alpha = ball_region_lit[:, :, 3:4] / 255.0
    rgb = ball_region_lit[:, :, :3]
    
    blended = bg_region * (1 - alpha) + rgb * alpha
    blended = blended.astype(np.uint8)
    
    # 放回背景
    result = background.copy()
    result[ball_y1:ball_y2, ball_x1:ball_x2] = blended
    
    bbox = (ball_x1, ball_y1, ball_x2, ball_y2)
    return result, bbox

def add_realistic_shadow(background, center, radius, shadow_intensity=0.4):
    """添加逼真的阴影"""
    shadow = np.zeros(background.shape[:2], dtype=np.float32)
    x, y = center
    
    # 创建椭圆阴影（稍微偏移）
    shadow_center = (int(x + radius * 0.3), int(y + radius * 0.3))
    axes = (int(radius * 0.8), int(radius * 0.3))
    
    cv2.ellipse(shadow, shadow_center, axes, 30, 0, 360, 1.0, -1)
    
    # 模糊阴影
    shadow = cv2.GaussianBlur(shadow, (25, 25), 10)
    
    # 应用阴影
    shadow_3d = np.repeat(shadow[..., np.newaxis], 3, axis=2)
    result = (background * (1 - shadow_3d * shadow_intensity)).astype(np.uint8)
    
    return result

def generate_natural_basketball_scene(background_path=None, texture_path=None, labels=None):
    """生成自然的篮球场景"""
    #background_size = (7680, 4320)
    background_size = (3840, 2160)
    # 创建背景
    if background_path:
        background = load_background_image(background_path, target_size=background_size)
    else:
        background = np.ones((480, 640, 3), dtype=np.uint8) * np.array([50, 120, 50], dtype=np.uint8)
        # 添加一些纹理
        cv2.rectangle(background, (0, 320), (640, 480), (80, 80, 80), 2)  # 地板
    
    # 创建篮球纹理
    if texture_path:
        basketball_texture = load_basketball_texture(texture_path)
    else:
        basketball_texture = create_spherical_basketball_texture()
    
    # 篮球位置和大小
    ball_position = labels[0][:2]
    ball_radius = int(labels[0][2]/2)
    
    print("=== ball radius ", ball_radius)
    # 先添加阴影
    background_with_shadow = add_realistic_shadow(
        background, ball_position, ball_radius, shadow_intensity=0.0
    )
    
    # 再添加篮球
    result, bbox = blend_basketball_naturally(
        background_with_shadow, basketball_texture, 
        ball_position, ball_radius
    )
    
    return result, bbox

def write_label(label_path, label):
    with open(label_path, 'w') as fout:
        fout.write('0 {}\n'.format(' '.join(str(_) for _ in label)))

if __name__ == '__main__':
    # 测试生成
    try:
        background_dir = '/media/avs/39efc736-bf77-4764-85e3-1a91d963a099/huangying/datas/三人篮球_4K_5_basket/clip_images/'
        label_dir = '/media/avs/39efc736-bf77-4764-85e3-1a91d963a099/huangying/datas/三人篮球_4K_5_basket/clip_images_labels'
        texture_dir = '/home/avs/Downloads/basketball_selected/'
        texture_imgs = os.listdir(texture_dir)
        label_backgrounds = os.listdir(label_dir)
        for idx in range(2000):
            label_name = label_backgrounds[idx % len(label_backgrounds)]
            label_json = os.path.join(label_dir, label_name)
            labels = load_label(label_json)
            name = label_name.split('.')[0]
            background_path = os.path.join(background_dir, '{}.jpg'.format(name))
            #texture_path = '/home/avs/Downloads/basketball_download/u=404125708,846798785&fm=253&fmt=auto&app=138&f=JPEG.jpeg'
            texture_path = os.path.join(texture_dir, texture_imgs[idx % len(texture_imgs)])

            result_image, bbox = generate_natural_basketball_scene(background_path, texture_path, labels)
            cv2.imwrite('basketball_generated/images/natural_basketball_{}_{}.jpg'.format(idx, name), result_image)
            write_label('basketball_generated/labels/natural_basketball_{}_{}.txt'.format(idx, name), labels[0][4:])

            print(f"成功生成图像！篮球边界框: {bbox}")
        
        # 显示结果
        #cv2.imshow('Natural Basketball', result_image)
        #cv2.waitKey(0)
        #cv2.destroyAllWindows()
        
    except Exception as e:
        print(f"错误: {e}")
        import traceback
        traceback.print_exc()
